from flask import Flask, request, jsonify
from flask_cors import CORS

from model.yiyan import YiyanChat
from model.deepseek import DeepseekChat
from model.chatglm import ChatglmChat
from model.llama import LlamaChat
from model.xunfei import XunfeiChat
from model.xuanyuan import XuanyuanChat

app = Flask(__name__)
CORS(app)  # 解决跨域问题

prompt = """
对于回复：先描述一下用户的需求，例如用户目前面临了什么问题，有什么需求，所以现在需要找法规和案例来支撑我的观点/诉求。
对于法条：请你详细给出相关规定，并写明性质、具体条文，并结合我的需求给出我可以如何使用这些具体规定。
对于案例：请你详细给出案例的裁判要旨、主要事实、法院观点等内容，使我可以准确了解了解该案件对我的价值。
"""

class ModelHandler:
    def __init__(self):
        self.yiyan_chat = YiyanChat()
        self.deepseek_chat = DeepseekChat()
        self.chatglm_chat = ChatglmChat()
        self.llama_chat = LlamaChat()
        self.xunfei_chat = XunfeiChat()
        self.xuanyuan_chat = XuanyuanChat()
        self.available_models = {
            "ERNIE-Bot-4": self.yiyan_chat,
            "DeepSeek-R1": self.deepseek_chat,
            "ChatGLM-4-Plus": self.chatglm_chat,
            "LLaMA-2": self.llama_chat,
            "Spark-4-Ultra": self.xunfei_chat,
            "XuanYuan": self.xuanyuan_chat,
        }

    # 处理前端请求
    def handle_request(self, data):
        model_type = data.get('model')
        user_message = data.get('message')
        
        if not self._validate_model(model_type):
            return {"error": "不支持的模型类型"}

        response = self._call_model(model_type, user_message + prompt)
        print(response)
        return response

    # 验证模型类型
    def _validate_model(self, model_type):
        return model_type in self.available_models

    # 调用模型
    def _call_model(self, model_type, message):
        if model_type not in self.available_models:
            raise ValueError(f"不支持的模型类型：{model_type}")
        return self.available_models[model_type].do_chat(message)

handler = ModelHandler()

@app.route('/chat', methods=['POST'])
def chat_endpoint():
    """聊天接口"""
    data = request.json
    result = handler.handle_request(data)
    return jsonify(result)

if __name__ == '__main__':
    app.run(port=5000, debug=True)
